#ifndef ag_avltree_avl_tree_h
#define ag_avltree_avl_tree_h

#include <vector>
#include <map>
#include <iostream>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <exception>


template<typename Key, typename Val>
class AVLTree {
public:

    //AVL Node
    /**
     * AVLNode struct
     * @param key
     * @param val
     */
    struct AvlNode {

        AvlNode(const Key& key, const Val& val) :
        parent_node(0), left_node(0), right_node(0), left_max_depth(0),
        right_max_depth(0), key(key), val(val) {
        }

        AvlNode* parent_node;
        AvlNode* left_node;
        AvlNode* right_node;
        int left_max_depth;
        int right_max_depth;
        Key key;
        Val val;
    };

    AvlNode* root_node;
    int size;
    bool doBalancing;

    //creates a new empty tree

    AVLTree() : root_node(0), size(0), doBalancing(true) {
    }

    //Destroys the tree
/**
 * Tree destructor
 */
    virtual ~AVLTree(void) {
        reset();
    }

    //reset (clear) the tree
/**
 * Clear the tree
 */
    void reset() {
        removeNode_(root_node);
    }

    //balance a node, invoke rotations on demand.
/**
 * balance a node, invoke rotations on demand.
 * @param parent_node
 * @return 
 */
    bool balanceNode_(AvlNode* parent_node) {
        bool didBalance = false;

        if (doBalancing) {
            int balancing = getBalance_(parent_node);

            if (balancing < -1) {
                int balanceRightHeavy = getBalance_(parent_node->left_node);
                if (balanceRightHeavy >= 1) {
                    slr(parent_node->left_node);
                    srr(parent_node);
                } else {
                    srr(parent_node);
                }
                didBalance = true;

            } else if (balancing > 1) {
                int balanceLeftHeavy = getBalance_(parent_node->right_node);
                if (balanceLeftHeavy <= -1) {
                    srr(parent_node->right_node);
                    slr(parent_node);
                } else {
                    slr(parent_node);
                }
                didBalance = true;

            }
        }

        return didBalance;
    }

    //remove an entry from the tree
/**
 * Remove an entry from the tree
 * @param key
 * @return 
 */
    bool remove(const Key& key) {
        AvlNode* node = getNode(key);
        if (!node) {
            return false;
        }

        this->size--;

        AvlNode* parent_node = node->parent_node;
        bool is_root_node = node->parent_node ? false : true;
        bool parent_pos_right = !is_root_node ? parent_node->right_node == node : false;
        bool has_left = node->left_node ? true : false;
        bool has_right = node->right_node ? true : false;

        //deleted node has no leaves
        if (!has_left && !has_right) {
            if (!is_root_node) {
                if (!parent_pos_right) {
                    parent_node->left_node = 0;
                    parent_node->left_max_depth--;
                } else {
                    parent_node->right_node = 0;
                    parent_node->right_max_depth--;
                }
            }

            delete node;

            if (is_root_node) {
                root_node = 0;
                return true;
            }
        }// deleted node has exactly one leaf
        else if ((has_left && !has_right) || (has_right && !has_left)) {
            bool detachRight = node->right_node ? true : false;

            if ((!parent_pos_right) && (!is_root_node)) {
                parent_node->left_node = detachRight ? node->right_node : node->left_node;
                parent_node->left_max_depth = (detachRight ? getMaxChildrenSize_(node->right_node) : getMaxChildrenSize_(node->left_node)) + 1;
            } else if (!is_root_node) {
                parent_node->right_node = detachRight ? node->right_node : node->left_node;
                parent_node->right_max_depth = (detachRight ? getMaxChildrenSize_(node->right_node) : getMaxChildrenSize_(node->left_node)) + 1;
            }

            if (detachRight) {
                node->right_node->parent_node = parent_node;
            } else {
                node->left_node->parent_node = parent_node;
            }

            if (is_root_node) {
                root_node = detachRight ? node->right_node : node->left_node;
                parent_node = root_node;
            }

            delete node;
        }// deleted node has 2 leaves
        else if (has_left && has_right) {
            AvlNode* replace_node = getLeftMostNode(node->right_node);
            AvlNode* propagation_root_node = replace_node->parent_node;
            node->key = replace_node->key;
            node->val = replace_node->val;

            if (replace_node != node->right_node) {
                replace_node->parent_node->left_node = 0;
                replace_node->parent_node->left_max_depth = 0;
                parent_node = propagation_root_node;
            } else {
                AvlNode* old_right_node = replace_node->right_node;
                node->right_node = old_right_node;
                node->right_max_depth = old_right_node ? (old_right_node->right_max_depth + 1) : 0;
                if (old_right_node) {
                    old_right_node->parent_node = node;
                }
                parent_node = node;
            }
            delete(replace_node);
        }

        bool didBalance = balanceNode_(parent_node);
        if (didBalance) {
            parent_node = parent_node->parent_node;
        }

        if (parent_node && parent_node->parent_node) {
            propagate_max_children_size(parent_node->parent_node, parent_node, getMaxChildrenSize_(parent_node), true);
        }
        return true;
    }

    //retrieves the leftmost child

    /**
     * Retrieve left most child
     * @param node
     * @return ptr to AVLNode of current node
     */
    AvlNode* getLeftMostNode(AvlNode* node) const {
        AvlNode* current_node = node;
        if (!current_node->left_node) {
            return current_node;
        } else {
            return getLeftMostNode(current_node->left_node);
        }
    }

    //dump the map into dotty format 

    /**
     * dump the map into dotty format 
     * @param sName
     */
    void dumpMap_(const std::string& sName) {
        std::ofstream ofs(sName.c_str());
        ofs << "digraph X { " << std::endl;
        drawNode_(ofs, root_node);
        ofs << "}" << std::endl;
    }

    /**
     * dump the map into dotty format 
     * @param sName filename
     */
    void outputInvertedIndex(ofstream &ofs) {
        //std::ofstream ofs(sName.c_str());
        outputInvertedNode_(ofs, root_node);
    }

    //Insert element into tree

    /**
     * Insert an element into the tree
     * @param key
     * @param val
     */
    void insert(const Key& key, const Val& val) {
        AvlNode* current_node = this->root_node;

        if (!current_node) {
            this->root_node = new AvlNode(key, val);
            size++;
            return;
        }

        for (;;) {
            if (current_node->key == key) {
                current_node->val = val;
                return;
            }

            if (current_node->key > key) {
                if (!current_node->left_node) {
                    current_node->left_node = newNode_(current_node, key, val);
                    size++;
                    propagate_max_children_size(current_node, current_node->left_node, 0, false);
                    return;
                }
                current_node = current_node->left_node;
            } else {
                if (!current_node->right_node) {
                    current_node->right_node = newNode_(current_node, key, val);
                    size++;
                    propagate_max_children_size(current_node, current_node->right_node, 0, false);
                    return;
                }
                current_node = current_node->right_node;
            }
        }
    }

    //Gets a key out of the avl tree

    /**
     * Gets a key out of the AVL Tree
     * @param key
     * @return 
     */
    Val search(const Key& key) const {
        AvlNode* ret = getNode(key);
        return ret ? ret->val : 0;
    }

    //Whether a key exists in the tree

    /**
     * Check if a key exists in the tree
     * @param key
     * @return true if it exists, false otherwise
     */
    bool contains(const Key& key) const {
        return getNode(key) ? true : false;
    }

    //Returns the size of the tree.

    /**
     * Get the size of the tree
     * @return size as int
     */
    int getSize() const {
        return this->size;
    }

protected:

    //Gets a key out of the avl tree

    /**
     * Gets a key out of the AVL Tree
     * @param key
     * @return 
     */
    AvlNode* getNode(const Key& key) const {
        AvlNode* current_node = this->root_node;
        if (!current_node) {
            return 0;
        }

        for (;;) {
            if (current_node->key == key) {
                return current_node;
            }

            if (current_node->key > key) {
                if (current_node->left_node) {
                    current_node = current_node->left_node;
                } else {
                    return 0;
                }
            } else {
                if (current_node->right_node) {
                    current_node = current_node->right_node;
                } else {
                    return 0;
                }
            }
        }
    }

    //Generates a new node (helper).

    /**
     * Generates a new node (helper)
     * @param parent
     * @param key
     * @param val
     * @return ptr to new Node
     */
    AvlNode* newNode_(AvlNode* parent, const Key& key, const Val& val) {
        AvlNode* newNode = new AvlNode(key, val);
        newNode->parent_node = parent;
        return newNode;
    }

    /**
     * Propagates new max children size to parents, does balancing on demand.
     * @param notified_node
     * @param sender_node
     * @param child_maxsize
     * @param is_deletion
     */
    void propagate_max_children_size(AvlNode* notified_node, AvlNode* sender_node, int child_maxsize, bool is_deletion) {
        bool isRight = sender_node == notified_node->right_node;
        int maxsize = child_maxsize + 1;
        int old_maxsize = getMaxChildrenSize_(notified_node);
        if (isRight) {
            notified_node->right_max_depth = maxsize;
        } else {
            notified_node->left_max_depth = maxsize;
        }
        int new_maxsize = getMaxChildrenSize_(notified_node);

        if (balanceNode_(notified_node)) {
            if (is_deletion) {
                notified_node = notified_node->parent_node; // our notified_node moved, readjust
            } else {
                return;
            }
        }

        if (notified_node->parent_node) {
            if (is_deletion ? new_maxsize != old_maxsize : new_maxsize > old_maxsize) {
                propagate_max_children_size(notified_node->parent_node, notified_node, new_maxsize, is_deletion);
            }
        }
    }

    //returns the balance for a given avl node 

    /**
     * returns the balance for a given avl node
     * @param node
     * @return 
     */
    int getBalance_(AvlNode* node) const {
        return node->right_max_depth - node->left_max_depth;
    }

    //returns the max children size

    /**
     * returns the max children size
     * @param node
     * @return 
     */
    int getMaxChildrenSize_(AvlNode* node) const {
        return std::max(node->left_max_depth, node->right_max_depth);
    }

    //remove Node and all of its children, deallocate memory, unlink the parents pointers (setting it to 0) and child count.

    /**
     * remove Node and all of its children, deallocate memory, unlink the parents pointers (setting it to 0) and child count.
     * @param node
     */
    void removeNode_(AvlNode* node) {
        if (!node) {
            return;
        }

        if (node->left_node) {
            removeNode_(node->left_node);
        }
        if (node->right_node) {
            removeNode_(node->right_node);
        }

        this->size--;

        if (node->parent_node) {
            if (node->parent_node->left_node == node) {
                node->parent_node->left_node = 0;
                node->parent_node->left_max_depth--;
            } else {
                node->parent_node->right_node = 0;
                node->parent_node->right_max_depth--;
            }
        } else {
            root_node = 0;
        }

        delete node;
    }

    //set tree balancing mode (useful for debugging purposes)

    /**
     * set tree balancing mode (useful for debugging purposes)
     * @param mode
     */
    void setBalancing(bool mode) {
        doBalancing = mode;
    }

    //single right rotate

    /**
     * Single right rotate
     * @param node
     */
    void srr(AvlNode* node) {
        bool newRoot = (node == root_node);
        node->left_node->parent_node = node->parent_node;
        AvlNode* old_right_node = node->left_node->right_node;
        node->left_node->right_node = node;
        node->parent_node = node->left_node;
        node->parent_node->right_node = node;
        node->left_node = old_right_node;
        node->left_max_depth = old_right_node ? getMaxChildrenSize_(old_right_node) : 0;
        if (old_right_node) {
            old_right_node->parent_node = node;
            node->left_max_depth++;
        }
        node->parent_node->right_max_depth++;

        if (newRoot) {
            root_node = node->parent_node;
        } else {
            if (node->parent_node->parent_node->left_node == node) {
                node->parent_node->parent_node->left_node = node->parent_node;
            } else {
                node->parent_node->parent_node->right_node = node->parent_node;
            }
        }
    }

    //single left rotate

    /**
     * Single left rotate
     * @param node
     */
    void slr(AvlNode* node) {
        bool newRoot = (node == root_node);
        node->right_node->parent_node = node->parent_node;
        AvlNode* old_left_node = node->right_node->left_node;
        node->right_node->left_node = node;
        node->parent_node = node->right_node;
        node->parent_node->left_node = node;
        node->right_node = old_left_node;
        node->right_max_depth = old_left_node ? getMaxChildrenSize_(old_left_node) : 0;
        if (old_left_node) {
            old_left_node->parent_node = node;
            node->right_max_depth++;
        }
        node->parent_node->left_max_depth++;

        if (newRoot) {
            root_node = node->parent_node;
        } else {
            if (node->parent_node->parent_node->left_node == node) {
                node->parent_node->parent_node->left_node = node->parent_node;
            } else {
                node->parent_node->parent_node->right_node = node->parent_node;
            }
        }
    }



    //create a node string representation 

    /**
     * create a node string representation
     * @param pCurr
     * @param ssOut
     */
    void createNode(AvlNode* pCurr, std::ostream& ssOut) {
        ssOut << "\"[k:" << pCurr->key << "], " << (int) (pCurr->right_max_depth - pCurr->left_max_depth) << "\"";
    }

    /**
     * create a node string representation
     * @param pCurr
     * @param ssOut
     */
    void createInvertedIndexNode(AvlNode* pCurr, std::ostream& ssOut) {
        ssOut << pCurr->key << "|";
        std::map<int, int>* documents = pCurr->val->getDocumentOccurances();

        for (std::map<int, int>::iterator docit = documents->begin(); docit != documents->end(); ++docit) {
            if (docit != documents->begin()) {
                ssOut << ',';
            }
            ssOut << docit->first;
        }
        ssOut << std::endl;
    }

    //draw a single node for dotty output

    /**
     * draw a single node for dotty output
     * @param fsOut
     * @param pCurr
     */
    void drawNode_(std::ofstream& fsOut, AvlNode* pCurr) {
        if (!pCurr) {
            return;
        }
        std::stringstream ssNode1, ssNode2, ssNode3, ssNodeTmp, ssNodeTmp2;
        createNode(pCurr, ssNode1);
        if (pCurr->left_node) {
            createNode(pCurr->left_node, ssNode2);
            createNode(pCurr->left_node->parent_node, ssNodeTmp);
            fsOut << ssNode1.str() << " -> " << ssNode2.str() << " [label=\"L\"];" << std::endl;
            fsOut << ssNode2.str() << " -> " << ssNodeTmp.str() << " [label=\"PL\"];" << std::endl;
        }
        if (pCurr->right_node) {
            createNode(pCurr->right_node, ssNode3);
            createNode(pCurr->right_node->parent_node, ssNodeTmp2);
            fsOut << ssNode1.str() << " -> " << ssNode3.str() << " [label=\"R\"];" << std::endl;
            fsOut << ssNode3.str() << " -> " << ssNodeTmp2.str() << " [label=\"PR\"];" << std::endl;
        }
        if (pCurr->left_node) {
            drawNode_(fsOut, pCurr->left_node);
        }
        if (pCurr->right_node) {
            drawNode_(fsOut, pCurr->right_node);
        }
        return;
    }

    /**
     * draw a single node for persistent inverted file index
     * @param fsOut
     * @param pCurr
     */
    void outputInvertedNode_(std::ofstream& fsOut, AvlNode* pCurr) {
        if (!pCurr) {
            return;
        }
        std::stringstream ssNode1, ssNode2, ssNode3, ssNodeTmp, ssNodeTmp2;
        createInvertedIndexNode(pCurr, ssNode1);
        if (pCurr->left_node) {
            createInvertedIndexNode(pCurr->left_node, ssNode2);
            createInvertedIndexNode(pCurr->left_node->parent_node, ssNodeTmp);
            fsOut << ssNode1.str();
            fsOut << ssNode2.str();
        }
        if (pCurr->right_node) {
            createInvertedIndexNode(pCurr->right_node, ssNode3);
            createInvertedIndexNode(pCurr->right_node->parent_node, ssNodeTmp2);
            fsOut << ssNode1.str();
            fsOut << ssNode3.str();
        }
        if (pCurr->left_node) {
            outputInvertedNode_(fsOut, pCurr->left_node);
        }
        if (pCurr->right_node) {
            outputInvertedNode_(fsOut, pCurr->right_node);
        }
        return;
    }



};
#endif
